#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "Wenbiao Xie"
__copyright__ = "Copyright 2023, The QNN Model convert&run wrap script project"
__version__ = "1.0.0"
__status__ = "Production"

import argparse
import os
import numpy as np

def to_shape(input_shape: str):
    if 'x' in input_shape:
        input_shape = input_shape.replace('x', ",")
    input_shape = input_shape.replace(' ', '')

    ss = input_shape.split(",")
    sizes = [int(x) for x in ss]
    sizes = tuple(sizes)
    return sizes

data_generator = None
def generate_tensor(size, dtype, min_val, max_val) -> np.ndarray[np.Any, np.dtype[np.float32]]:
    global data_generator
    if data_generator is None:
        data_generator = np.random.default_rng()

    print(f"start to generate tensor shape={size}, dtype={dtype}, ranges=[{min_val}, {max_val})...")
    data_array = data_generator.random(size, dtype=np.float32, out=None)
    if dtype == np.float32:
        if min_val != 0 or max_val != 1.0:
            diff = max_val - min_val
            data_array = data_array * diff + min_val

    elif dtype == np.float16:
        if min_val != 0 or max_val != 1.0:
            diff = max_val - min_val
            data_array = data_array * diff + min_val
        data_array = data_array.astype(dtype)
    else:
        diff = max_val - min_val
        data_array = data_array * diff + min_val
        data_array = np.ceil(data_array).astype(dtype)

    return data_array

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-s","--shape",  dest='input_shape', default=None, type=str, help="the shape of target tensor, such as WxHxC or W,H,C")
    parser.add_argument("-t","--dtype",  dest='input_dtype', 
                        default="float32", 
                        choices=["float16", "float32", "uint8", "uint16", "uint32"],
                        type=str, 
                        help="the dtype of target tensor")
    parser.add_argument("--min", default=0.0, type=float, help="the min value of target tensor")
    parser.add_argument("--max", default=1.0, type=float, help="the max value of target tensor")
    parser.add_argument("-o", "--output",  dest='output_file', default=None, type=str, help="the output file path of target tensor")

    options = parser.parse_args()
    if options.input_shape is None or len(options.input_shape) == 0:
        raise ValueError("invalid input_shape!")
    
    if options.min >= options.max:
        raise ValueError("invalid min and max parameter, should be max > min!")    

    sizes = to_shape(options.input_shape)
    dtype = np.dtype(options.input_dtype)
    
    ranges_array = np.array([options.min, options.max], dtype=np.float32)
    if dtype != np.float32:
        ranges_array = ranges_array.astype(dtype)
    min_val = ranges_array[0]
    max_val = ranges_array[1]
    data_array = generate_tensor(sizes, dtype, min_val, max_val)
    print(f"save tensor to {options.output_file}...")
    basedir = os.path.dirname(options.output_file)
    os.makedirs(basedir, exist_ok=True)    
    data_array.tofile(options.output_file)
    print(f"saved successfully!")


    
